import numpy as np

manual_seed =42

sigma_value=0.2

np.random.seed(manual_seed)             # 设置 NumPy 随机数种子
class ParallelGraph:
    def __init__(self, theta_x, theta_y, num_of_rounds, k):
        self.type = 'parallel'
        self.theta_x = theta_x
        self.theta_y = theta_y
        self.n = len(theta_x) + 2
        self.num_parents_y = self.n - 2
        self.delta = 1 / self.n / np.sqrt(num_of_rounds)
        self.delta=self.delta
        self.T = num_of_rounds
        self.k = k

    def compute_rho_lr(self, t):
    
        return (np.sqrt(self.n * np.log(1 + self.n * t) + 2 * np.log(1 / self.delta)) + np.sqrt(self.n)) / 10

    def compute_rho_ofu(self):
        return 3 * np.sqrt(np.log(3 / self.delta)) / 10

    def simulate_explore(self, intervention):
        x = [np.random.random() < self.theta_x[i] for i in range(self.num_parents_y)]
        for i in intervention:
            x[int(i)] = 0.0
        noise = np.random.normal(0, 100, 1)
        arr_x=np.array(x)
        y_vec=np.array(x) * self.theta_y
        y=  np.sum(y_vec)+noise
        # y = (np.random.random() < np.sum(np.array(x) * self.theta_y))
        return   np.array([x]).T, y

    def simulate_cluster_y(self, intervention):
      

        # np.random.seed(manual_seed)             # 设置 NumPy 随机数种子
        # x = [np.random.normal(0, sigma_value, 1)[0] * self.theta_x[i] for i in range(self.num_parents_y)]
        x = [np.random.random() < self.theta_x[i] for i in range(self.num_parents_y)]


        for i in intervention:
            x[int(i)] = 1.0
        noise = np.random.normal(0, sigma_value, 1)
        arr_x=np.array(x)
       

       
        y_vec=np.array(x) * self.theta_y
        y=  np.sum(y_vec)+noise

    
        y_g_one= np.sum(y_vec[:3])
        y_g_two=np.sum(y_vec[-3:])

       
 
        # y = (np.random.random() < np.sum(np.array(x) * self.theta_y))
        return   np.array([x]).T, y,   y_g_one, y_g_two

    def simulate(self, intervention):
    

        # np.random.seed(manual_seed)             # 设置 NumPy 随机数种子
        # x = [np.random.normal(0, sigma_value, 1)[0] * self.theta_x[i] for i in range(self.num_parents_y)]
        x = [np.random.random() < self.theta_x[i] for i in range(self.num_parents_y)]
      
        for i in intervention:
            x[int(i)] = 1
        noise = np.random.normal(0,  sigma_value, 1)
        y_vec=np.array(x) * self.theta_y
      

        y=np.sum(np.array(x) * self.theta_y)+ noise
 
        # y = (np.random.random() < np.sum(np.array(x) * self.theta_y))
        return np.array([x]).T, y

    def simulate_component(self, intervention):
    

        # np.random.seed(manual_seed)             # 设置 NumPy 随机数种子
        # x = [np.random.normal(0, sigma_value, 1)[0] * self.theta_x[i] for i in range(self.num_parents_y)]
        x = [np.random.random() < self.theta_x[i] for i in range(self.num_parents_y)]
      
        for i in intervention:
            x[int(i)] = 1
        noise = np.random.normal(0,  sigma_value, 1)
        y_vec=np.array(x) * self.theta_y

        y_list=[]

        for j in range(6):
            y_now= y_vec[j]+noise
          
        
            y_list.append( y_now)

      
   
        # y=np.sum(np.array(x) * self.theta_y)+ noise
 
        # y = (np.random.random() < np.sum(np.array(x) * self.theta_y))
        return np.array([x]).T, y_list

    def simulate_two(self, intervention):
    

        # np.random.seed(manual_seed)             # 设置 NumPy 随机数种子
        # x = [np.random.normal(0, sigma_value, 1)[0] * self.theta_x[i] for i in range(self.num_parents_y)]
        x = [np.random.random() < self.theta_x[i] for i in range(self.num_parents_y)]
      
        for i in intervention:
            x[int(i)] = 1
        noise = np.random.normal(0,  sigma_value, 1)
        y_vec=np.array(x) * self.theta_y

        y_list=[]

        
        y_one=np.sum(np.array(x)[:2] * self.theta_y[:2])+ noise
        y_list.append( y_one)

        y_two=np.sum(np.array(x)[-2:] * self.theta_y[-2:])+ noise
        y_list.append(  y_two)

        y_three=np.sum(np.array(x)[2:4] * self.theta_y[2:4])+ noise
        y_list.append(  y_three)

      
   
        # y=np.sum(np.array(x) * self.theta_y)+ noise
 
        # y = (np.random.random() < np.sum(np.array(x) * self.theta_y))
        return np.array([x]).T, y_list


    def expect_y(self, intervention):
        x = self.theta_x.copy()
        for i in intervention:
            x[int(i)] = 1
        return np.sum(np.array(x) * self.theta_y)

    def best_expect_y(self):
        max_y = - 9999
        for j in range(np.power(self.num_parents_y, self.k)):
            intervened_indexes = []
            for k in range(self.k):
                index = int(j / np.power(self.num_parents_y, k)) % self.num_parents_y
                if index not in intervened_indexes:
                    intervened_indexes.append(index)
            if len(intervened_indexes) < self.k:
                continue
            expect_y = self.expect_y(intervened_indexes)
            if expect_y > max_y:
                max_y = expect_y
        return max_y
